#include "reid.h"

#include <opencv2/opencv.hpp>

#include <chrono>
#include <fstream>

using namespace perception::camera;
using namespace perception::camera;

std::shared_ptr<ReID> ReID::detector_ = nullptr;

ReID::ReID()
    : logger__(rclcpp::get_logger("ReID"))
{
}

bool ReID::init(const std::string& model_dir)
{
    const std::string engine_file_path = model_dir + "/reid.engine";
    RCLCPP_INFO(logger__, "engine file: %s", engine_file_path.c_str());
    std::ifstream model_file(engine_file_path, std::ios::binary);
    if (!model_file.good()) {
        RCLCPP_FATAL(logger__, "model file not exist:%s", engine_file_path.c_str());
    }

    model_file.seekg(0, std::ios::end);
    auto size = model_file.tellg();
    model_file.seekg(0, std::ios::beg);
    char* trtModelStream = new char[size];
    assert(trtModelStream);
    model_file.read(trtModelStream, size);
    model_file.close();
    initLibNvInferPlugins(&this->logger_, "");
    this->runtime_ = nvinfer1::createInferRuntime(this->logger_);
    assert(this->runtime_ != nullptr);

    this->engine_ = this->runtime_->deserializeCudaEngine(trtModelStream, size);
    assert(this->engine_ != nullptr);
    delete[] trtModelStream;
    this->context_ = this->engine_->createExecutionContext();

    assert(this->context_ != nullptr);
    cudaStreamCreate(&this->stream_);
    this->num_bindings = this->engine_->getNbBindings();

    for (int i = 0; i < this->num_bindings; ++i) {
        ReIDBinding binding;
        nvinfer1::Dims dims;
        nvinfer1::DataType dtype = this->engine_->getBindingDataType(i);
        std::string name = this->engine_->getBindingName(i);
        binding.name = name;
        binding.dsize = type_to_size(dtype);

        bool IsInput = engine_->bindingIsInput(i);
        if (IsInput) {
            this->num_inputs += 1;
            dims = this->engine_->getProfileDimensions(
                i,
                0,
                nvinfer1::OptProfileSelector::kMAX);
            binding.size = get_size_by_dims(dims);
            binding.dims = dims;
            this->input_bindings.push_back(binding);
            // set max opt shape
            this->context_->setBindingDimensions(i, dims);

        } else {
            dims = this->context_->getBindingDimensions(i);
            binding.size = get_size_by_dims(dims);
            binding.dims = dims;
            this->output_bindings.push_back(binding);
            this->num_outputs += 1;
        }
        std::cout << "output binding:" << name << ",dsize:" << binding.dsize << ",size:" << binding.size << ",dims:" << binding.dims.d[0] << std::endl;
    }
    data_ = new (float[128 * 128 * 3]);
    make_pipe(true);
    return true;
}

void ReID::make_pipe(bool warmup)
{

    for (auto& bindings : this->input_bindings) {
        void* d_ptr { nullptr };
        cudaError_t en = cudaMalloc(
            &d_ptr,
            bindings.size * bindings.dsize);
        RCLCPP_INFO(logger__, "errno:%d", en);
        this->device_ptrs.push_back(d_ptr);
    }

    for (auto& bindings : this->output_bindings) {
        void *d_ptr { nullptr }, *h_ptr { nullptr };
        // size_t size = 29*8400*sizeof(float);//640 * 640 * 3 * sizeof(float);//
        size_t size = bindings.size * bindings.dsize;
        // size:512,dsize:4
        RCLCPP_INFO(logger__, "output bindings size:%ld,dsize:%ld", bindings.size, bindings.dsize);
        cudaError_t en = cudaMalloc(
            &d_ptr,
            size);
        RCLCPP_INFO(logger__, "cudama:%d", en);
        en = cudaHostAlloc(
            &h_ptr,
            size,
            0);
        RCLCPP_INFO(logger__, "host:%d", en);
        this->device_ptrs.push_back(d_ptr);
        this->host_ptrs.push_back(h_ptr);
    }

    if (warmup) {
        for (int i = 0; i < 10; i++) {
            for (auto& bindings : this->input_bindings) {
                // size_t size = 640 * 640 * 3 * sizeof(float);//
                size_t size = bindings.size * bindings.dsize;
                void* h_ptr = malloc(size);
                memset(h_ptr, 0, size);
                // cudaError_t en = 
                cudaMemcpy(
                    this->device_ptrs[0],
                    h_ptr,
                    size,
                    cudaMemcpyHostToDevice);
                free(h_ptr);
            }
            this->infer();
        }
        RCLCPP_INFO(logger__, "model warmup 10 times2");
    }
}

void ReID::deinit()
{
    this->context_->destroy();
    this->engine_->destroy();
    this->runtime_->destroy();
    cudaStreamDestroy(this->stream_);
    delete[] data_;
    for (auto& ptr : this->device_ptrs) {
        assert(cudaFree(ptr));
    }

    for (auto& ptr : this->host_ptrs) {
        assert(cudaFreeHost(ptr));
    }

    detector_ = nullptr;
    return;
}

void ReID::letterbox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
{
    const float inp_h = size.height;
    const float inp_w = size.width;
    float height = image.rows;
    float width = image.cols;

    float r = std::min(inp_h / height, inp_w / width);
    int padw = std::round(width * r);
    int padh = std::round(height * r);

    cv::Mat tmp;
    if ((int)width != padw || (int)height != padh) {
        cv::resize(image, tmp, cv::Size(padw, padh));
    } else {
        tmp = image.clone();
    }

    float dw = inp_w - padw;
    float dh = inp_h - padh;

    dw /= 2.0f;
    dh /= 2.0f;
    int top = int(std::round(dh - 0.1f));
    int bottom = int(std::round(dh + 0.1f));
    int left = int(std::round(dw - 0.1f));
    int right = int(std::round(dw + 0.1f));

    cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, { 114, 114, 114 });
    // std::cout << "resize wh:" << tmp.cols << "," << tmp.rows << std::endl;
    // uint8_t [] buf =  new uint8_t[tmp.cols * tmp.rows * 3];
    uint8_t* src = (uint8_t*)tmp.data;
    int imh = tmp.rows;
    int imw = tmp.cols;
    out = cv::Mat::zeros(imh, imw, CV_32FC3);
    float* buf = (float*)out.data;
    for (int i = 0; i < imh; ++i) {
        for (int j = 0; j < imw; ++j) {
            // for(int c = 0; c < 3; ++c){
            //     // mean: 0.485, 0.456, 0.406, std: 0.229, 0.224, 0.225
            //     buf[c * imh * imw + i * imw + j] = src[i * imw * c + j * 3 + c] / 255. - ;
            // }
            // rgb, mean: 0.485, 0.456, 0.406, std: 0.229, 0.224, 0.225
            // rgb <= bgr

            buf[2 * imh * imw + i * imw + j] = (src[i * imw * 3 + j * 3 + 0] / 255. - 0.406) / 0.225;
            buf[1 * imh * imw + i * imw + j] = (src[i * imw * 3 + j * 3 + 1] / 255. - 0.456) / 0.224;
            buf[0 * imh * imw + i * imw + j] = (src[i * imw * 3 + j * 3 + 2] / 255. - 0.485) / 0.229;
        }
    }

    // cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);
    // this->pparam.ratio  = 1 / r;
    // this->pparam.dw     = dw;
    // this->pparam.dh     = dh;
    // this->pparam.height = height;
    // this->pparam.width  = width;
}
// 只支持1920x1080图片，不是这个尺寸会resize
void ReID::copy_from_mat(const cv::Mat& image)
{
    // auto time_cpy_start = std::chrono::high_resolution_clock::now();
    cv::Mat blob;
    // auto tt0 = std::chrono::system_clock::now();
    // int infer_h = 360, infer_w = 640;
    image_orig_ = image.clone();

    // cv::resize(image, blob, { infer_w_, infer_h_ }, 0, 0, cv::INTER_NEAREST);
    // image_resize_ = blob.clone();
    auto size = cv::Size(infer_w_, infer_h_);
    letterbox(image, blob, size);
    blob.convertTo(blob, CV_32FC3);
    // cv::cvtColor(blob, blob, cv::COLOR_BGR2RGB);
    // auto time_resize_end = std::chrono::high_resolution_clock::now();
    this->context_->setBindingDimensions(
        0,
        nvinfer1::Dims { 4,
            { 1, 3, infer_h_, infer_w_ } });

    cudaMemcpy(
        this->device_ptrs[0],
        blob.data,
        infer_w_ * infer_h_ * 3 * sizeof(float),
        cudaMemcpyHostToDevice);
    // auto time_cpy_end = std::chrono::high_resolution_clock::now();
    // auto time_resize = std::chrono::duration_cast<std::chrono::microseconds>(time_resize_end - time_cpy_start).count();
    // auto time_cpy = std::chrono::duration_cast<std::chrono::microseconds>(time_cpy_end - time_resize_end).count();
    // std::cout << "reid copy from mat resize time:" << time_resize << ", mem cpy:" << time_cpy << std::endl;
}

void ReID::infer()
{
    // auto time_infer_start = std::chrono::high_resolution_clock::now();
    this->context_->enqueueV2(
        this->device_ptrs.data(),
        this->stream_,
        nullptr);
    // auto time_equeue_end = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < this->num_outputs; i++) {
        // size_t osize = 29*8400*sizeof(float);//640 * 640 * 3 * sizeof(float);//
        size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
        cudaMemcpy(this->host_ptrs[i],
            this->device_ptrs[i + this->num_inputs],
            osize,
            cudaMemcpyDeviceToHost);
    }
    // auto time_cpy_end = std::chrono::high_resolution_clock::now();
    cudaStreamSynchronize(this->stream_);
    // auto time_sync_end = std::chrono::high_resolution_clock::now();
    // auto time_enqeue = std::chrono::duration_cast<std::chrono::microseconds>(time_equeue_end - time_infer_start).count();
    // auto time_cpy = std::chrono::duration_cast<std::chrono::microseconds>(time_cpy_end - time_equeue_end).count();
    // auto time_sync = std::chrono::duration_cast<std::chrono::microseconds>(time_sync_end - time_cpy_end).count();
    // std::cout << "reid infer enqeue:" << time_enqeue << ", mem cpy:" << time_cpy << ", sync:" << time_sync << std::endl;
}

void ReID::postprocess(std::vector<float>& feat)
{
    // auto time_start = std::chrono::high_resolution_clock::now();
    // std::cout << "outbinding size:" <<this->output_bindings.size() << std::endl;
    int feat_dim = this->output_bindings[0].dims.d[1];
    // std::cout << "feat_dim:" << feat_dim << std::endl;
    // std::cout << "this->host_ptrs size:" << this->host_ptrs.size() << std::endl;
    float* hfeat = static_cast<float*>(this->host_ptrs[0]); // 1x1
    float power_sum = 0;
    for (int i = 0; i < feat_dim; ++i) {
        feat.push_back(hfeat[i]);
        power_sum += hfeat[i] * hfeat[i];
        // std::cout << "feat:" << i << ":" << hfeat[i] << std::endl;
    }
    power_sum = std::sqrt(power_sum);
    for (int i = 0; i < feat_dim; ++i) {
        feat[i] /= power_sum;
        power_sum += hfeat[i] * hfeat[i];
    }
    // std::cout << std::endl;
    power_sum = std::sqrt(power_sum);
    float sum = 0;
    for (int i = 0; i < feat_dim; ++i) {
        feat[i] /= power_sum;
        sum += feat[i];
        // std::cout << "feat1 " << i << ":" << feat[i];
    }
    // std::cout << "power_sum:" << power_sum << std::endl;
    // auto time_end = std::chrono::high_resolution_clock::now();
    // auto time_tr  = std::chrono::duration_cast<std::chrono::microseconds>(time_end - time_start).count();
    // std::cout << "reid post use:" << time_tr << std::endl;
}

int ReID::process(const cv::Mat image, std::vector<float>& feat)
{
    if (image.empty()) {
        std::cout << "reid image empty!!! return false" << std::endl;
        return false;
    }

    // float score_thres = 0.5; // 0.25f;
    {
        ++infer_count_;
        // std::lock_guard<std::mutex> guard(pipe_mutex_);
        auto start_time = std::chrono::high_resolution_clock::now();
        cv::Mat image_clone = image.clone();
        copy_from_mat(image_clone);
        infer();
        postprocess(feat);
        // postprocess_cpu(objs, score_thres);
        auto end_time = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
        time_total_ += duration.count();
        // std::cout << "reid total:" << time_total_ << ", count:" << infer_count_ << std::endl;
        // std::cout << "reid process time: " << 1.0 * time_total_ / infer_count_ << ", cur:" << duration.count()  << std::endl;
    }
    return 0;
}
float ReID::cosine_dist(const std::vector<float>& normed_feat1, const std::vector<float>& normed_feat2)
{
    size_t len = normed_feat1.size();
    assert(len == normed_feat2.size());
    float sum = 0;
    for (size_t i = 0; i < len; ++i){
        sum += normed_feat1[i] * normed_feat2[i];
    }
    float dist = 1 - sum;
    return dist;
}
std::shared_ptr<ReID> ReID::get_detector(const std::string& model_dir)
{
    if (!detector_) {
        detector_ = std::shared_ptr<ReID>(new ReID());
        detector_->init(model_dir);
    }
    return detector_;
}
